
import pandas as pd
import numpy as np
from itertools import chain
import pickle
import time
import networkx as nx
from walker import RandomWalker
from sklearn.preprocessing import LabelEncoder
import argparse
import collections
import copy


def get_mashup_top_type(top_num):
    mashup_info = pd.read_csv(args.data_path + 'mashup_nodes_estimator.csv', sep="\t")
    mashup_info = mashup_info[['name', 'c']]
    type_num = collections.Counter(mashup_info['c']).most_common(top_num)
    print(type_num)
    top_type = [item[0] for item in type_num]
    mashup_info = mashup_info[mashup_info['c'].isin(top_type)]

    return set(mashup_info['name'])

def get_api_type():
    api_data = pd.read_csv(args.data_path + 'api_nodes_estimator.csv', sep='\t')
    api_type = sorted(list(set(api_data['c'])))
    print(api_type)




def get_graph_context_all_pairs(walks, window_size):
    all_pairs = []
    for k in range(len(walks)):
        for i in range(len(walks[k])):
            for j in range(i - window_size, i + window_size + 1):
                if i == j or j < 0 or j >= len(walks[k]):
                    continue
                else:
                    all_pairs.append([walks[k][i], walks[k][j]])
    return np.array(all_pairs, dtype=np.int32)



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='manual to this script')
    parser.add_argument("--data_path", type=str, default='./data/')
    parser.add_argument("--p", type=float, default=0.25)
    parser.add_argument("--q", type=float, default=2)
    parser.add_argument("--num_walks", type=int, default=10)
    parser.add_argument("--walk_length", type=int, default=10)
    parser.add_argument("--window_size", type=int, default=5)
    args = parser.parse_known_args()[0]



    source = []
    target = []
    action_data = pd.read_csv(args.data_path + 'm-a_edges.csv', sep = '\t')

    get_api_type()
    #挑选top——tyep的mashup
    top_mashup = get_mashup_top_type(30)
    print(top_mashup)

    print(action_data)
    action_data = action_data[action_data['source'].isin(top_mashup)]
    print(action_data)

    #类别编码
    all_api = action_data['target'].unique()
    api_lbe = LabelEncoder()
    api_lbe.fit_transform(all_api)
    action_data['target'] = api_lbe.transform(action_data['target'])

    mashup_lbe = LabelEncoder()
    mashup_lbe.fit_transform(list(top_mashup))
    action_data['source'] = mashup_lbe.transform(action_data['source'])



    # add side info
    api_data = pd.read_csv(args.data_path + 'api_nodes_estimator.csv', sep='\t')
    api_data = api_data[['url', 'c', 'oac', 'ac']]
    print(api_data)

    all_api = pd.DataFrame({'url': list(all_api)})

    api_side_info = pd.merge(all_api, api_data, on='url', how='left').fillna('null')
    api_side_info['oac'] = api_side_info['oac'].apply(str)
    api_side_info['ac'] = api_side_info['ac'].apply(str)
    print(api_side_info)

    # id2index
    for feat in api_side_info.columns:
        if feat != 'url':
            lbe = LabelEncoder()
            api_side_info[feat] = lbe.fit_transform(api_side_info[feat])
        else:
            api_side_info[feat] = api_lbe.transform(api_side_info[feat])
    print(api_side_info)

    api_side_info = api_side_info.sort_values(by=['url'], ascending=True)
    api_side_info.to_csv('./data_cache/api_side_info.csv', index=False, header=False, sep='\t')
    api_side_info = api_side_info.rename(columns={'url': 'target'})

    mashup_data = pd.read_csv(args.data_path + 'mashup_nodes_estimator.csv', sep='\t')
    mashup_data = mashup_data[['name', 'c', 'oac', 'ac']]

    all_mashup = pd.DataFrame({'name': list(top_mashup)})

    mashup_side_info = pd.merge(all_mashup, mashup_data, on='name', how='left').fillna('null')
    mashup_side_info['oac'] = mashup_side_info['oac'].apply(str)
    mashup_side_info['ac'] = mashup_side_info['ac'].apply(str)
    print(mashup_side_info)

    # id2index
    for feat in mashup_side_info.columns:
        if feat != 'name':
            lbe = LabelEncoder()
            mashup_side_info[feat] = lbe.fit_transform(mashup_side_info[feat])
        else:
            mashup_side_info[feat] = mashup_lbe.transform(mashup_side_info[feat])
    mashup_side_info = mashup_side_info.rename(columns={'name': 'source'})
    print(mashup_side_info)

    #构造deepFM数据
    print('deepFM----------------------------------------------------------------------------------------------')
    agg_data = action_data.groupby('source').agg(list)
    apilist = set(api_side_info['target'])
    def make_neg_api(x):
        neg_list = list(apilist.difference(x.target))
        return np.random.choice(neg_list, 10)
    agg_data['fake_target'] = agg_data.apply(make_neg_api, axis=  1)
    print(agg_data)


    def make_neg_data(df):
        columns = ['source', 'target']
        '''创建一个DataFrame，即我们需要还原的qq_name'''
        neg_data = pd.DataFrame(columns=columns)
        '''遍历索引'''
        for each in df.index:
            api_list = df.loc[each, 'fake_target']
            for api in api_list:
                neg_data = neg_data.append(pd.DataFrame([[each, api]], columns=columns))

        return neg_data

    neg_data = make_neg_data(agg_data)
    print(neg_data)

    pos_data = copy.copy(action_data)
    pos_data['label'] = 1
    neg_data['label'] = 0
    deepFM_data = pd.concat([pos_data,neg_data])
    deepFM_data = pd.merge(deepFM_data, api_side_info, how='left', on='target')
    deepFM_data = pd.merge(deepFM_data, mashup_side_info, how='left', on='source')
    deepFM_data = deepFM_data.sample(frac=1)

    print(deepFM_data)
    deepFM_data.to_csv('./data_cache/deepFM_data.csv', sep='\t', index= False)

    #划分EGSE数据集
    print('EGSE--------------------------------------------------------------------------')
    EGES_test_data = action_data.sample(frac=0.2)

    EGES_action_data = action_data.drop(labels=EGES_test_data.index)


    EGES_action_data = action_data.groupby('source').agg(list)
    print(EGES_action_data)
    print(EGES_action_data.columns)

    EGES_test_data = EGES_test_data.groupby('source').agg(list)
    print(EGES_test_data)
    print(EGES_test_data.columns)

    EGES_test_data.to_csv('./data_cache/test_data.csv', sep='\t')
    EGES_action_data.to_csv('./data_cache/train_data.csv', sep='\t')


    #  利用共现关系构建关联
    action_data = action_data.groupby('source').agg(list)
    session_list = action_data['target']

    start_time = time.time()
    session_list_all = []
    for session in session_list:
            if len(session) > 1:
                session_list_all.append(session)
    print('make session list done, time cost {0}'.format(str(time.time() - start_time)))

    # session2graph
    node_pair = dict()
    for session in session_list_all:
        for i in range(0, len(session) - 1 ):
            for j in range(i + 1, len(session)):
                if session[i] != session[j]:
                    if (session[i], session[j]) not in node_pair.keys():
                        node_pair[(session[i], session[j])] = 1
                        node_pair[(session[j], session[i])] = 1
                    else:
                        node_pair[(session[i], session[j])] += 1
                        node_pair[(session[j], session[i])] += 1

    in_node_list = list(map(lambda x: x[0], list(node_pair.keys())))
    out_node_list = list(map(lambda x: x[1], list(node_pair.keys())))
    weight_list = list(node_pair.values())
    graph_df = pd.DataFrame({'in_node': in_node_list, 'out_node': out_node_list, 'weight': weight_list})
    graph_df.to_csv('./data_cache/graph.csv', sep=' ', index=False, header=False)

    G = nx.read_edgelist('./data_cache/graph.csv', create_using=nx.DiGraph(), nodetype=None, data=[('weight', int)])
    walker = RandomWalker(G, p=args.p, q=args.q)
    print("Preprocess transition probs...")
    walker.preprocess_transition_probs()

    session_reproduce = walker.simulate_walks(num_walks=args.num_walks, walk_length=args.walk_length, workers=4,
                                              verbose=1)
    session_reproduce = list(filter(lambda x: len(x) > 1, session_reproduce))

    # get pair
    all_pairs = get_graph_context_all_pairs(session_reproduce, args.window_size)
    np.savetxt('./data_cache/all_pairs', X=all_pairs, fmt="%d", delimiter=" ")



